#include <stdlib.h>
#include <assert.h>
#include <stdio.h>
#include <sys/time.h>
#include <string.h>

#include "net.h"
#include "quark.h"
#include "model.h"

flow_model_t *_flow_model_create(flow_t *f1, flow_t *f2) {
  flow_model_t *model;

  model = calloc(1, sizeof(*model));
  assert(model);

  model->protocol = f1->protocol;
  model->references[0] = f1;
  model->references[1] = f2;
  TAILQ_INIT(&model->packets);

  return model;
}

static int _sockaddr_in_cmp(const struct sockaddr_in *a,
  const struct sockaddr_in *b) {
  if (a->sin_addr.s_addr < b->sin_addr.s_addr)
    return -1;
  if (a->sin_addr.s_addr > b->sin_addr.s_addr)
    return 1;
  if (a->sin_port < b->sin_port)
    return -1;
  if (a->sin_port > b->sin_port)
    return 1;
  return 0;
}

flowset_model_t *flowset_pair_flows(flowset_t *set1, flowset_t *set2) {
  flowset_model_t *model;
  flow_t *f1, *f2;
  flow_model_t *fm;
  struct timeval *t1, *t2;
  int may_have_secondary = 0;
  flow_t *missing_f1, *missing_f2;

  model = calloc(1, sizeof(*model));
  assert(model);

  model->references[0] = set1;
  model->references[1] = set2;
  TAILQ_INIT(&model->flows);
  TAILQ_INIT(&model->packets);
  TAILQ_INIT(&model->variables);

  t1 = &set1->flows.tqh_first->when;
  t2 = &set2->flows.tqh_first->when;

  /* first step: tie primary flows */
  TAILQ_FOREACH(f1,&set1->flows,next) {
    TAILQ_FOREACH(f2,&set2->flows,next) {
      /* same protocol */
      if (f1->protocol != f2->protocol)
        continue;
      /* same type: dialog or monolog */
      if (f1->type != f2->type)
        continue;
      /* at least same connector or same connectee */
      if ((_sockaddr_in_cmp(&f1->connector,&f2->connector)) &&
        (_sockaddr_in_cmp(&f1->connectee,&f2->connectee)))
        continue;
      break;
    }

    if (!f2) {
      may_have_secondary++;
      continue;
    }

    fm = _flow_model_create(f1, f2);
    timersub(&f1->when,t1,&fm->relative_start);
    if (_sockaddr_in_cmp(&f1->connector,&f2->connector) == 0)
      memcpy(&fm->connector, &f1->connector, sizeof(f1->connector));
    else if (f1->connector.sin_addr.s_addr == f2->connector.sin_addr.s_addr) {
      fm->connector.sin_addr.s_addr = f2->connector.sin_addr.s_addr;
    }
    else {}
    if (_sockaddr_in_cmp(&f1->connectee,&f2->connectee) == 0)
      memcpy(&fm->connectee, &f1->connectee, sizeof(f1->connectee));
    else if (f1->connectee.sin_addr.s_addr == f2->connectee.sin_addr.s_addr) {
      fm->connectee.sin_addr.s_addr = f2->connectee.sin_addr.s_addr;
    }
    else {}
    fm->type = PRIMARY_FLOW;
    TAILQ_INSERT_TAIL(&model->flows,fm,next);

    model->n_flows++;
  }

  /* RTP fits here as udp port is not constant and is given a SIP message */
  /* workaround: if there is one flow not classified yet in each set, then match them together */
  if (may_have_secondary && model->n_flows == set1->n_flows-1 && model->n_flows == set2->n_flows-1) {
    missing_f1 = missing_f2 = NULL;
    TAILQ_FOREACH(f1,&set1->flows,next) {
      TAILQ_FOREACH(fm,&model->flows,next) {
        if (fm->references[0] == f1)
          break;
      }
      /* fm == NULL means that f1 is not part model->flows, ie the missing flow */
      if (!fm) {
        missing_f1 = f1;
        break;
      }
    }
    TAILQ_FOREACH(f2,&set2->flows,next) {
      TAILQ_FOREACH(fm,&model->flows,next) {
        if (fm->references[1] == f2)
          break;
      }
      if (!fm) {
        missing_f2 = f2;
        break;
      }
    }
    assert(missing_f1 && missing_f2);
    fm = _flow_model_create(missing_f1, missing_f2);
    timersub(&f1->when,t1,&fm->relative_start);
    if (missing_f1->connector.sin_addr.s_addr == missing_f2->connector.sin_addr.s_addr)
      fm->connector.sin_addr.s_addr = missing_f1->connector.sin_addr.s_addr;
    if (missing_f1->connectee.sin_addr.s_addr == missing_f2->connectee.sin_addr.s_addr)
      fm->connectee.sin_addr.s_addr = missing_f1->connectee.sin_addr.s_addr;
    fm->type = SECONDARY_FLOW;
    TAILQ_INSERT_TAIL(&model->flows,fm,next);

    model->n_flows++;
  }
  else if (may_have_secondary) {
    printf("way too complicated for me, good bye!\n");
    abort();
  }

  return model;
}

static chunk_t *_constant_chunk_create(quark_t *q1, quark_t *q2) {
  chunk_t *v;

  v = calloc(1, sizeof(*v));
  assert(v);

  v->type = CONSTANT;
  v->_u.constant_quark.segment.iov_base = malloc(q1->segment.iov_len);
  assert(v->_u.constant_quark.segment.iov_base);
  memcpy(v->_u.constant_quark.segment.iov_base, q1->segment.iov_base, q1->segment.iov_len);
  v->_u.constant_quark.segment.iov_len = q1->segment.iov_len;

  return v;
}

static chunk_t *_variable_chunk_create(quark_t *q1, quark_t *q2) {
  chunk_t *v;

  v = calloc(1, sizeof(*v));
  assert(v);

  v->type = VARIABLE;
  v->_u.variable_quark.references[0] = q1;
  v->_u.variable_quark.references[1] = q2;
  assert(q1 && q2);

  return v;
}

static packet_model_t *_packet_model_create(chunked_packet_t *pkt1,
  chunked_packet_t *pkt2) {
  packet_model_t *pm;
  quark_t *q1, *q2;
  chunk_t *v;

  pm = calloc(1, sizeof(*pm));
  assert(pm);
  TAILQ_INIT(&pm->chunks);
  pm->references[0] = pkt1;
  pm->references[1] = pkt2;

  if (pkt1->n_quarks != pkt2->n_quarks) {
    printf("trouble arising, there is not the same number of quarks between the two packets!\n");
    return pm;
  }

  for (q1 = TAILQ_FIRST(&pkt1->quarks), q2 = TAILQ_FIRST(&pkt2->quarks);
    q1 && q2; q1 = TAILQ_NEXT(q1,next_by_packet), q2 = TAILQ_NEXT(q2,next_by_packet)) {
    if (q1->syntax != q2->syntax) {
      printf("trouble ahead, two possibly similar quarks do not have the same syntax!\n");
      continue;
    }
    if (quark_cmp(q1, q2) == 0) {
      v = _constant_chunk_create(q1, q2);
    }
    else {
      v = _variable_chunk_create(q1, q2);
    }

    TAILQ_INSERT_TAIL(&pm->chunks,v,next);
  }

  return pm;
}

static packet_model_t *_packet_model_lookup(flowset_model_t *model,
  packet_t *pkt) {
  flow_model_t *fm;
  packet_model_t *pm;
  TAILQ_FOREACH(fm,&model->flows,next) {
    TAILQ_FOREACH(pm,&fm->packets,next_by_flow) {
      if (pm->references[0]->packet == pkt)
        return pm;
      if (pm->references[1]->packet == pkt)
        return pm;
    }
  }
  return NULL;
}

static variable_t *_flowset_model_lookup(flowset_model_t *model,
  quark_t *quark) {
  variable_t *v;
  TAILQ_FOREACH(v,&model->variables,next) {
    if (quark_cmp(v->references[0], quark) == 0)
      return v;
  }
  return NULL;
}

static variable_t *_flowset_model_create(flowset_model_t *model,
  quark_t *quarks[2]) {
  variable_t *v;

  v = calloc(1, sizeof(*v));
  assert(v);

  v->references[0] = quarks[0];
  v->references[1] = quarks[1];

  TAILQ_INSERT_TAIL(&model->variables,v,next);

  return v;
}

void flowset_analyze(flowset_model_t *model, struct chunked_packets *pkts1,
  struct chunked_packets *pkts2) {
  flow_model_t *fm;
  packet_model_t *pm;
  chunked_packet_t *pkt1, *pkt2;
  packet_t *p;
  chunk_t *c;
  variable_t *var;
  int id = 0;
  char lport[16];
  char rport[16];
  int ret;
  quark_t lportq, rportq;
  chunk_t *lportc, *rportc;
  quark_t *lquarks[2];
  quark_t *rquarks[2];

  TAILQ_FOREACH(fm,&model->flows,next) {
    pkt1 = chunked_packets_by_flow(pkts1, fm->references[0]);
    pkt2 = chunked_packets_by_flow(pkts2, fm->references[1]);

    if (fm->references[0]->type == MONOLOG_FLOW) {
      printf("treating flow model %p as static\n", fm);
      continue;
    }

    for (; pkt1 && pkt2; pkt1 = TAILQ_NEXT(pkt1,next_by_flow),
      pkt2 = TAILQ_NEXT(pkt2,next_by_flow)) {
      if (pkt1->packet->type != pkt2->packet->type) {
        printf("there is a packet inversion between the two captures, can't deal with this!\n");
        abort();
      }

      pm = _packet_model_create(pkt1, pkt2);
      pm->flow = fm;
      timersub(&pm->references[0]->packet->when,&fm->references[0]->when,&pm->relative_start);
      TAILQ_INSERT_TAIL(&fm->packets,pm,next_by_flow);
    }
  }

  TAILQ_FOREACH(p,&model->references[0]->packets,next_by_time) {
    pm = _packet_model_lookup(model, p);
    if (!pm)
      continue;
    TAILQ_INSERT_TAIL(&model->packets,pm,next_by_time);
  }

  /* flow endpoints lookup */
  /* just port for the moment */
  TAILQ_FOREACH(fm,&model->flows,next) {
    /* we do not look for all occurences, just the very first one */
    lportc = NULL;
    rportc = NULL;
    ret = snprintf(lport, sizeof(lport)-1, "%d", ntohs(fm->references[0]->connector.sin_port));
    lportq.segment.iov_base = lport;
    lportq.segment.iov_len = ret;
    lportq.next_by_packet.tqe_next = NULL;
    ret = snprintf(rport, sizeof(rport)-1, "%d", ntohs(fm->references[0]->connectee.sin_port));
    rportq.segment.iov_base = rport;
    rportq.segment.iov_len = ret;
    rportq.next_by_packet.tqe_next = NULL;
    printf("flow l4 lookup: %d %d\n", ntohs(fm->references[0]->connector.sin_port), ntohs(fm->references[0]->connectee.sin_port));
    for (pm = TAILQ_FIRST(&model->packets); pm && timercmp(&pm->references[0]->packet->when,&fm->references[0]->when,<);
      pm = TAILQ_NEXT(pm,next_by_time)) {
      /* no self reference */
      if (pm->references[0]->packet->flow == fm->references[0])
        continue;
      TAILQ_FOREACH(c,&pm->chunks,next) {
        switch (c->type) {
        case CONSTANT:
          if (quark_cmp(&lportq, &c->_u.constant_quark) == 0 && !lportc) {
            lportc = c;
            lquarks[0] = &c->_u.constant_quark;
            lquarks[1] = NULL;
          }
          if (quark_cmp(&rportq, &c->_u.constant_quark) == 0 && !rportc) {
            rportc = c;
            rquarks[0] = &c->_u.constant_quark;
            rquarks[1] = NULL;
          }
          break;
        case VARIABLE:
          if (quark_cmp(&lportq, c->_u.variable_quark.references[0]) == 0 && !lportc) {
            lportc = c;
            lquarks[0] = c->_u.variable_quark.references[0];
            lquarks[1] = c->_u.variable_quark.references[1];
          }
          if (quark_cmp(&rportq, c->_u.variable_quark.references[0]) == 0 && !rportc) {
            rportc = c;
            rquarks[0] = c->_u.variable_quark.references[0];
            rquarks[1] = c->_u.variable_quark.references[1];
          }
          break;
        default:
          break;
        }
      }
      if (rportc && lportc)
        break;
    }
    /* if something match, turn it into a variable */
    if (rportc) {
      var = _flowset_model_lookup(model, rquarks[0]);
      if (!var) {
        var = _flowset_model_create(model, rquarks);
        var->id = id++;
      }
      fm->connectee_port = var;
      rportc->type = INSTANCE;
      rportc->_u.instance = var;
    }
    if (lportc) {
      var = _flowset_model_lookup(model, lquarks[0]);
      if (!var) {
        var = _flowset_model_create(model, lquarks);
        var->id = id++;
      }
      fm->connector_port = var;
      lportc->type = INSTANCE;
      lportc->_u.instance = var;
    }
  }

  /* perform equality tests between variable chunks */
  TAILQ_FOREACH(pm,&model->packets,next_by_time) {
    TAILQ_FOREACH(c,&pm->chunks,next) {
      if (c->type == CONSTANT)
        continue;
      if (c->type == INSTANCE)
        continue;
      var = _flowset_model_lookup(model, c->_u.variable_quark.references[0]);
      if (!var) {
        var = _flowset_model_create(model, c->_u.variable_quark.references);
        var->id = id++;
      }
      c->type = INSTANCE;
      c->_u.instance = var;
    }
  }
}
